
#include <errno.h>
#include <libintl.h>
#include <stdlib.h>
#include <unistd.h>
#include <ldsodefs.h>

#include "dynamic-link.h"
#include <sysdep.h>

#include "dl-load-i.h"
#include "dl-load-pe.h"

#include <rewine.h>

#include <not-cancel.h>

int
is_pe_magic (const char *buf)
{
  return (buf[0] == 'M'
       && buf[1] == 'Z') ? 1 : 0;
}

int
do_verify_pe (int fd, struct filebuf *fbp, struct link_map *loader,
              int mode, bool *found_other_class,
              int *perrval, const char **perrstr)
{
  return 0;
}

static struct link_map *peloader = NULL;

#define REWINE_FUNC(name) \
  typedef typeof (rewine_##name) *pfn##name##_t; \
  static pfn##name##_t pfn##name;

#define SET_REWINE_FUNC(name) \
  pfn##name = (pfn##name##_t) do_sym (peloader, "rewine_" #name);

REWINE_FUNC (AddFindPath);
REWINE_FUNC (OpenImageFd);
REWINE_FUNC (CloseImage);
REWINE_FUNC (GetImageBase);
REWINE_FUNC (GetImageSize);
REWINE_FUNC (UpdateELFSymtab);
REWINE_FUNC (GetELFInfo);

// cannot use _dl_sym, for that leads to a multi-defined error
static void *
do_sym(struct link_map *map, const char *name) {
  int flags = DL_LOOKUP_RETURN_NEWEST;
  const ElfW(Sym) *ref = NULL;

  lookup_t result = GLRO(dl_lookup_symbol_one) (
    name, NULL /* undef_map */,
    map, NULL /* version */, 0 /* type_class */, flags,
    &ref
  );
  if (ref != NULL) {
    //_dl_error_printf("%s found %lx %lx\n", name, ref->st_value, (intptr_t) DL_SYMBOL_ADDRESS (result, ref));
    return DL_SYMBOL_ADDRESS (result, ref);
  }
  _dl_error_printf("%s not found\n", name);
  return NULL;
}

// we have put it in deps chain, so just find it out
static int
init_peloader (void) {
  if (!peloader) {
    peloader = _dl_find_loaded_map (LM_ID_BASE, "librewine.so");
    if (!peloader) {
      _dl_error_printf ("peloader open failed\n");
      return -1;
    }
    SET_REWINE_FUNC (AddFindPath);
    SET_REWINE_FUNC (OpenImageFd);
    SET_REWINE_FUNC (CloseImage);
    SET_REWINE_FUNC (GetImageBase);
    SET_REWINE_FUNC (GetImageSize);
    SET_REWINE_FUNC (UpdateELFSymtab);
    SET_REWINE_FUNC (GetELFInfo);

    pfnAddFindPath (".");
  }
  return 0;
}

static int
init_dynamic_section (
    struct link_map *l, ElfW(Addr) needed, ElfW(Addr) soname,
    ElfW(Addr) hash, ElfW(Addr) symtab, ElfW(Addr) strtab, ElfW(Xword) strsz);

static int
load_dll (struct link_map *l, int fd, int dynamic, int *perrval, const char **perrstr)
{
  rw_image_t handle = pfnOpenImageFd (fd);
  if (!handle) {
    *perrstr = N_("rewine_OpenImage failed");
    goto call_lose_errno;
  call_lose_errno:
    *perrval = errno;
  call_lose:
    return 1;
  }
  l->l_dll_handle = handle;

  if (!dynamic) {
    // link-load
    ElfW(Sym) *symtab = (ElfW(Sym) *) D_PTR (l, l_info[DT_SYMTAB]);
    char *strtab = (char *) D_PTR (l, l_info[DT_STRTAB]);
    // we do a trick here for 'symcnt'
    pfnUpdateELFSymtab(handle, (uint32_t) symtab[0].st_size, symtab, strtab);
  } else {
    // dynamic-load
    rw_elf_t info = pfnGetELFInfo (handle, l->l_libname->name);
    int ret = init_dynamic_section (l, (ElfW(Addr)) info.needed, (ElfW(Addr)) info.soname,
        (ElfW(Addr)) info.hash, (ElfW(Addr)) info.symtab, (ElfW(Addr)) info.strtab, (ElfW(Xword)) info.strsz);
    if (ret) {
      *perrstr = N_("init .dynamic failed");
      *perrval = 0;
      goto call_lose;
    }
  }

  l->l_map_start = (ElfW(Addr)) pfnGetImageBase(handle);
  l->l_map_end = (ElfW(Addr)) ((char *)l->l_map_start + pfnGetImageSize(handle));

  // we do not need to use ELF reloc
  //l->l_addr = l->l_map_start;

  return 0;
}

static void myinit (struct link_map *, int, char **, char **);
static void myfini (struct link_map *);

static int
init_dynamic_section (
    struct link_map *l, ElfW(Addr) needed, ElfW(Addr) soname,
    ElfW(Addr) hash, ElfW(Addr) symtab, ElfW(Addr) strtab, ElfW(Xword) strsz)
{
  l->l_ldnum = 10; // include trailing-NULL
  l->l_ld = (ElfW(Dyn) *) malloc (sizeof (ElfW(Dyn)) * l->l_ldnum);
  if (!l->l_ld) {
    return -1;
  }
  ElfW(Dyn) *ld = l->l_ld;
  ld->d_tag = DT_NEEDED;
  ld->d_un.d_ptr = needed;
  ++ld;
  ld->d_tag = DT_SONAME;
  ld->d_un.d_ptr = soname;
  ++ld;
  ld->d_tag = DT_INIT;
  ld->d_un.d_ptr = (ElfW(Addr)) &myinit;
  ++ld;
  ld->d_tag = DT_FINI;
  ld->d_un.d_ptr = (ElfW(Addr)) &myfini;
  ++ld;
  ld->d_tag = DT_GNU_HASH;
  ld->d_un.d_ptr = hash;
  ++ld;
  ld->d_tag = DT_STRTAB;
  ld->d_un.d_ptr = strtab;
  ++ld;
  ld->d_tag = DT_SYMTAB;
  ld->d_un.d_ptr = symtab;
  ++ld;
  ld->d_tag = DT_STRSZ;
  ld->d_un.d_val = strsz;
  ++ld;
  ld->d_tag = DT_SYMENT;
  ld->d_un.d_val = sizeof (ElfW(Sym));
  ++ld;
  ld->d_tag = DT_NULL;
  ld->d_un.d_ptr = (ElfW(Addr)) NULL;
  ++ld;

  elf_get_dynamic_info (l, NULL);
  _dl_setup_hash (l);
  return 0;
}

typedef struct rw_elf_hash_t {
  uint32_t nbuckets;
  uint32_t symbias;
  uint32_t nbitmask; // must be a power of 2.
  uint32_t shift2;
  uint32_t data[0];
  // bitmask[__ELF_NATIVE_CLASS/32 * nbitmask]
  // buckets[nbuckets]
  // hashes[ndefines]
} hash_t;

typedef struct rw_elf_strtab_t {
  uint32_t count;
  size_t size;
  char *ptr;
  ElfW(Addr) needed;
  ElfW(Addr) soname;
} rw_elf_strtab_t;

static ElfW(Addr)
append_strtab (rw_elf_strtab_t *strtab, const char *s, size_t len)
{
  if (!s) return 0;
  if (!len) len = strlen (s);
  strtab->count++;
  ElfW(Addr) oldsize = (ElfW(Addr)) strtab->size;
  strtab->size += len + 1;
  char *ptr;
  if (!strtab->ptr) {
    ptr = (char *) malloc (strtab->size);
  } else {
    ptr = (char *) realloc (strtab->ptr, strtab->size);
  }
  if (!ptr) {
    _dl_error_printf ("realloc failed: %s\n", s);
    free (strtab->ptr);
    strtab->ptr = NULL;
    return 0;
  } else {
    strtab->ptr = ptr;
    memcpy (ptr + oldsize, s, len + 1);
    return oldsize;
  }
}

static rw_elf_strtab_t
init_elf_strtab (const char *soname)
{
  rw_elf_strtab_t strtab;
  strtab.count = 0;
  strtab.size = 0;
  strtab.ptr = NULL;
  append_strtab (&strtab, "", 0);
  strtab.needed = append_strtab (&strtab, "librewine.so", 0);
  strtab.soname = append_strtab (&strtab, soname, 0);
  return strtab;
}

static uint_fast32_t
dl_new_hash (const char *s)
{
  uint_fast32_t h = 5381;
  for (unsigned char c = *s; c != '\0'; c = *++s)
    h = h * 33 + c;
  return h & 0xffffffff;
}

static hash_t *
init_hash (ElfW(Word) ndef, ElfW(Addr) symtab, ElfW(Addr) strtab)
{
  uint32_t nbuckets = 1;
  uint32_t cbitmask = __ELF_NATIVE_CLASS / 32;
  uint32_t nbitmask = 2;
  hash_t *hash = (hash_t *) malloc(sizeof(hash_t) + sizeof(uint32_t) * (cbitmask * nbitmask + nbuckets + ndef));
  hash->nbuckets = nbuckets;
  hash->symbias = 1; // for symtab[0]
  hash->nbitmask = nbitmask;
  hash->shift2 = 1;

  uint32_t *pdata = &hash->data[0];
  // filters
  for (int i = 0; i < cbitmask * nbitmask; ++i) {
    *pdata++ = -1; // 100% pass
  }
  // buckets
  *pdata++ = 1; // start at symtab[symbias]

  // hashes
  ElfW(Sym) *psym = (ElfW(Sym) *) symtab+1 + ndef-1;
  pdata += ndef-1;
  for (int i = ndef; i; --i, --psym, --pdata) {
    const char *name = (char *) strtab + psym->st_name;
    uint_fast32_t hash = dl_new_hash (name) & ~1;
    uint32_t lsb = (i == ndef) ? 1 : 0; // || (hash != (*(pdata+1) & ~1)) ? 1 : 0;
    *pdata = hash | lsb;
    //_dl_error_printf ("%lx %x: %s\n", hash, *pdata, name);
  }
  return hash;
}

static int
fake_dll (struct link_map *l, int fd, int *perrval, const char **perrstr)
{
  const char *soname = l->l_libname->name;

  size_t szRead;
  IMAGE_DOS_HEADER dos;
  __lseek (fd, (off_t) 0, SEEK_SET);
  szRead = sizeof (IMAGE_DOS_HEADER);
  if ((size_t) __read_nocancel (fd, (void *)&dos, szRead) != szRead) {
    _dl_error_printf ("%s: broken dos header\n", soname);
    return 1;
  }
  __lseek (fd, (off_t) dos.e_lfanew, SEEK_SET);
  IMAGE_NT_HEADERS nt;
  memset (&nt, 0, sizeof (nt));
  szRead = sizeof (nt.Signature) + sizeof (nt.FileHeader);
  if ((size_t) __read_nocancel (fd, (void *)&nt, szRead) != szRead) {
    _dl_error_printf ("%s: broken nt header\n", soname);
    return 1;
  }
  szRead = nt.FileHeader.SizeOfOptionalHeader;
  if ((size_t) __read_nocancel (fd, (void *)&nt.OptionalMagic, szRead) != szRead) {
    _dl_error_printf ("%s: broken optonal header\n", soname);
    return 1;
  }
  IMAGE_SECTION_HEADER hdrs[nt.FileHeader.NumberOfSections];
  szRead = sizeof (hdrs);
  if ((size_t) __read_nocancel (fd, (void *)&hdrs[0], szRead) != szRead) {
    _dl_error_printf ("%s: broken section headers\n", soname);
    return 1;
  }

  //intptr_t imagebase = nt.OptionalNt.ImageBase;

  if (IMAGE_DIRECTORY_ENTRY_EXPORT >= nt.OptionalNt.NumberOfRvaAndSizes) {
    _dl_error_printf ("%s: no export dentry\n", soname);
    return 1;
  }
  PIMAGE_DATA_DIRECTORY dentry = &nt.OptionalNt.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
  if (!dentry->Size) {
    _dl_error_printf ("%s: no export table (size=0)\n", soname);
    return 1;
  }
  intptr_t tblstart = dentry->VirtualAddress;
  intptr_t tblend = tblstart + dentry->Size;

  char exportd_raw[dentry->Size];
  PIMAGE_EXPORT_DIRECTORY exportd = NULL;

  PIMAGE_SECTION_HEADER hdr = &hdrs[0];
  for (int i = 0; i < nt.FileHeader.NumberOfSections; ++i, ++hdr) {
    intptr_t secstart = hdr->VirtualAddress;
    intptr_t secend = secstart + hdr->Misc.VirtualSize;
    //_dl_error_printf ("%s: %lx-%lx : %s %lx-%lx\n", soname, tblstart, tblend, hdr->Name, secstart, secend);
    if (secstart <= tblstart && tblend <= secend) {
      __lseek (fd, (off_t) hdr->PointerToRawData + (tblstart - secstart), SEEK_SET);
      szRead = dentry->Size;
      if ((size_t) __read_nocancel (fd, (void *) exportd_raw, szRead) != szRead) {
        _dl_error_printf ("%s: broken export table\n", soname);
        return 1;
      }
      exportd = (PIMAGE_EXPORT_DIRECTORY) exportd_raw;
      break;
    }
  }
  if (!exportd) {
    _dl_error_printf ("%s: export table not found\n", soname);
    return 1;
  }

  intptr_t rvadiff = tblstart;

  const char *dllname = (char *) (exportd_raw + (exportd->Name - rvadiff));
  //_dl_error_printf ("DLL: %s\n", dllname);

  ElfW(Word) nsym = exportd->NumberOfNames;
  ElfW(Addr) symtab = (ElfW(Addr)) malloc (sizeof (ElfW(Sym)) * (1 + nsym));

  ElfW(Sym) *psym = (ElfW(Sym) *) symtab;
  // symtab starts with NULL symbol
  psym->st_name = 0;
  psym->st_value = (ElfW(Addr)) NULL;
  psym->st_size = nsym; // for nowhere else can pass 'nsym' to librewine
  psym->st_info = ELFW(ST_INFO) (STB_LOCAL, STT_NOTYPE);
  psym->st_other = 0;
  psym->st_shndx = 0;
  ++psym;

  // realloc() must be called consistently
  rw_elf_strtab_t strtab = init_elf_strtab (dllname);

  PDWORD tblNames = (PDWORD) (exportd_raw + (exportd->AddressOfNames - rvadiff));
  PWORD tblOrdinals = (PWORD) (exportd_raw + (exportd->AddressOfNameOrdinals - rvadiff));
  PDWORD tblValues = (PDWORD) (exportd_raw + (exportd->AddressOfFunctions - rvadiff));
  for (int i = 0; i < exportd->NumberOfNames; ++i, ++tblNames, ++tblOrdinals, ++psym) {
    const char *name = (char *) (exportd_raw + (*tblNames - rvadiff));
    WORD ordinal = *tblOrdinals;
    DWORD value = tblValues[ordinal];
    //_dl_error_printf ("#%u %s: %x\n", ordinal, name, value);

    psym->st_name = append_strtab (&strtab, name, strlen(name));
    psym->st_value = (ElfW(Addr)) value;
    psym->st_size = sizeof (void *);
    psym->st_info = ELFW(ST_INFO) (STB_GLOBAL, STT_NOTYPE);
    psym->st_other = STV_DEFAULT;
    psym->st_shndx = 1; // anything but SHN_ABS
  }

  ElfW(Addr) hashtbl = (ElfW(Addr)) init_hash (nsym, symtab, (ElfW(Addr)) strtab.ptr);

  int ret = init_dynamic_section (l, (ElfW(Addr)) strtab.needed, (ElfW(Addr)) strtab.soname,
      (ElfW(Addr)) hashtbl, (ElfW(Addr)) symtab, (ElfW(Addr)) strtab.ptr, (ElfW(Xword)) strtab.size);

  if (ret) {
    _dl_error_printf ("%s: init .dynamic failed: %d\n", soname, ret);
    return 1;
  }
  return 0;
}

typedef struct lazyload_t {
  struct lazyload_t *next;
  int fd;
  Lmid_t nsid;
  struct link_map *map;
} lazyload_t;

static lazyload_t m_load_head; // placeholder
static lazyload_t *m_load_tail = &m_load_head;

int
do_load_pe (int fd, struct filebuf *fbp,
            struct link_map *l, struct link_map *loader, Lmid_t nsid,
            int mode, void **stack_endp,
            int *perrval, const char **perrstr)
{
  l->l_dll = 1;

  // no need to init loadable sections
  l->l_phnum = 0;
  l->l_phdr = NULL;

  l->l_map_start = (ElfW(Addr)) NULL;
  l->l_map_end = (ElfW(Addr)) ((char *)l->l_map_start + 0);
  l->l_contiguous = 1;

  // va = rva + l_addr
  l->l_addr = 0;

  // not support symbol versioning
  l->l_versyms = NULL;
  l->l_nversions = 0;
  l->l_versions = NULL;
  
  // no need to setup TLS
  l->l_need_tls_init = 0;

  // no need to reloc
  l->l_relro_addr = (ElfW(Addr)) NULL;
  l->l_relro_size = 0;

  // no need to call entrypoint
  l->l_entry = 0;

  lazyload_t *entry = (lazyload_t *) malloc(sizeof(lazyload_t));
  entry->next = NULL;
  entry->fd = fd;
  entry->nsid = nsid;
  entry->map = l;

  m_load_tail->next = entry;
  m_load_tail = entry;

  if (!peloader) {
    return fake_dll (l, fd, perrval, perrstr);
  } else {
    return load_dll (l, fd, 1 /* dynamic-load */, perrval, perrstr);
  }
}

typedef struct lazyreloc_t {
  struct lazyreloc_t *next;
  ElfW(Addr) *r_addr;
  const ElfW(Rela) *reloc;
  unsigned long int r_type;
  struct link_map *resolved_map;
  const ElfW(Sym) *resolved_sym;
} lazyreloc_t;

typedef struct lazyreloc_map_t {
  struct lazyreloc_map_t *next;
  struct link_map *map;
  struct lazyreloc_t head; // placeholder
  struct lazyreloc_t *tail;
} lazyreloc_map_t;

static lazyreloc_map_t m_reloc_head; // placeholder
static lazyreloc_map_t *m_reloc_tail = &m_reloc_head;

void add_lazyreloc (
    struct link_map *reloc_map, ElfW(Addr) *const r_addr, const ElfW(Rela) *reloc, unsigned long int r_type,
    struct link_map *resolved_map, const ElfW(Sym) *resolved_sym)
{
  lazyreloc_t *entry = (lazyreloc_t *) malloc (sizeof (lazyreloc_t));
  entry->next = NULL;
  entry->r_addr = r_addr;
  entry->reloc = reloc;
  entry->r_type = r_type;
  entry->resolved_map = resolved_map;
  entry->resolved_sym = resolved_sym;

  lazyreloc_map_t *map_entry = m_reloc_head.next;
  while (map_entry) {
    if (map_entry->map == reloc_map) break;
    map_entry = map_entry->next;
  }
  if (!map_entry) {
    map_entry = (lazyreloc_map_t *) malloc (sizeof (lazyreloc_map_t));
    map_entry->next = NULL;
    map_entry->map = reloc_map;
    map_entry->head.next = NULL;
    map_entry->tail = &map_entry->head;

    m_reloc_tail->next = map_entry;
    m_reloc_tail = map_entry;
  }
  map_entry->tail->next = entry;
  map_entry->tail = entry;
}

static void
myinit (struct link_map *l, int argc, char *argv[], char *envp[])
{
  if (peloader) return;

  init_peloader ();

  lazyload_t *loadhead = m_load_head.next;
  lazyload_t *loaditem = loadhead;
  m_load_head.next = NULL;
  m_load_tail = &m_load_head;

  //lazyload_t *lastload = NULL;
  while (loaditem) {
    int errval;
    const char *errstr;
    int ret = load_dll (loaditem->map, loaditem->fd, 0 /* linkload */, &errval, &errstr);
    if (ret) {
      __close_nocancel (loaditem->fd);
      _dl_error_printf ("load %s failed: %s\n", loaditem->map->l_libname->name, errstr);
      //TODO: call_lose
    }

    //lastload = loaditem;
    loaditem = loaditem->next;
    //free (lastload);
  }

  lazyreloc_map_t *reloc_map_head = m_reloc_head.next;
  lazyreloc_map_t *reloc_map_item = reloc_map_head;
  m_reloc_head.next = NULL;
  m_reloc_tail = &m_reloc_head;

  struct textrels {
    caddr_t start;
    size_t len;
    int prot;
    struct textrels *next;
  } *textrels = NULL;
  const char *errstring = NULL;

  //lazyreloc_map_t *last_reloc_map = NULL;
  while (reloc_map_item) {
    struct link_map *l = reloc_map_item->map;
    const ElfW(Phdr) *ph;
    for (ph = l->l_phdr; ph < &l->l_phdr[l->l_phnum]; ++ph) {
      if (ph->p_type == PT_LOAD) {
        struct textrels *newp = (struct textrels *) alloca (sizeof (*newp));
        newp->len = ALIGN_UP (ph->p_vaddr + ph->p_memsz, GLRO(dl_pagesize)) - ALIGN_DOWN (ph->p_vaddr, GLRO(dl_pagesize));
        newp->start = PTR_ALIGN_DOWN (ph->p_vaddr, GLRO(dl_pagesize)) + (caddr_t) l->l_addr;

        newp->prot = 0;
        if (ph->p_flags & PF_R)
          newp->prot |= PROT_READ;
        if (ph->p_flags & PF_W)
          newp->prot |= PROT_WRITE;
        if (ph->p_flags & PF_X)
          newp->prot |= PROT_EXEC;

        if (__mprotect (newp->start, newp->len, newp->prot|PROT_WRITE) < 0) {
          errstring = N_("cannot make segment writable for relocation");
        call_error:
          _dl_signal_error (errno, l->l_name, NULL, errstring);
        }

        newp->next = textrels;
        textrels = newp;
      }
    }

    lazyreloc_t *relocitem = reloc_map_item->head.next;
    reloc_map_item->head.next = NULL;
    reloc_map_item->tail = &reloc_map_item->head;

    //lazyreloc_t *lastreloc = NULL;
    while (relocitem) {
      ElfW(Addr) *r_addr = relocitem->r_addr;
      const ElfW(Rela) *reloc = relocitem->reloc;
      ElfW(Addr) value = SYMBOL_ADDRESS (relocitem->resolved_map, relocitem->resolved_sym, true);

      __attribute__((unused)) 
      const char *strtab = (const void *) D_PTR (relocitem->resolved_map, l_info[DT_STRTAB]);
      _dl_printf ("reloc symbol %s with value %lx\n", strtab + relocitem->resolved_sym->st_name, value);

      switch (relocitem->r_type) {
        case R_X86_64_GLOB_DAT:
        case R_X86_64_JUMP_SLOT:
          *(r_addr) = value + reloc->r_addend;
          break;
        case R_X86_64_64:
          *((Elf64_Addr *) r_addr) = (Elf64_Addr) value + reloc->r_addend;
          break;
        case R_X86_64_32:
          *((Elf32_Addr *) r_addr) = value + reloc->r_addend;
          break;
        case R_X86_64_PC32:
          //TODO: overflow detection
          *((Elf32_Addr *) r_addr) = value + reloc->r_addend - (ElfW(Addr)) r_addr;
          break;
        case R_X86_64_COPY:
          //TODO        
          break;
        case R_X86_64_IRELATIVE:
          value = l->l_addr + reloc->r_addend;
          value = ((ElfW(Addr) (*) (void)) value) ();
          *(r_addr) = value;
          break;
        default:
          _dl_reloc_bad_type (l, relocitem->r_type, 0);
          break;
      }

      //lastreloc = relocitem;
      relocitem = relocitem->next;
      //free (lastreloc);
    }

    while (textrels) {
      if (__mprotect (textrels->start, textrels->len, textrels->prot) < 0) {
        errstring = N_("cannot restore segment prot after reloc");
        goto call_error;
      }
      textrels = textrels->next;
    }

    if (l->l_relro_size != 0)
      _dl_protect_relro (l);

    //last_reloc_map = reloc_map_item;
    reloc_map_item = reloc_map_item->next;
    //free (last_reloc_map);
  }
}

static void
myfini (struct link_map *l)
{
  if (!peloader || !l->l_dll_handle) return;
  pfnCloseImage (l->l_dll_handle);
}
